
#include <apr_general.h>
#include <apr_getopt.h>
#include <apr_strings.h>
#include "iface/APRInterface.h"
#include "iface/CryptInterface.h"
#include "iface/SQLiteInterface.h"
#include "iface/XMLInterface.h"

#include "db/KeyDB.h"
#include "crypt/RSAKeyPair.h"
#include "net/TCPClient.h"

#include "net/Broker4.h"
#include "utils/BaseConvert.h"
#include "utils/Hash.h"
#include "utils/StringUtils.h"
#include "utils/SerialBlockSink.h"
#include "utils/SerialBlockSource.h"
#include "utils/RunningWorkProcess.h"

#include "net/Broker4Client.h"
#include "net/BrokerMessage.h"

#include "xml/XMLDescriptor.h"

#include "default_hosts.h"

#define COMPUTENAME ".compute"

typedef struct compute_args
{
   char *brokername;
   char *hostname;
   char *hostport;
   char *xmlfile;
   char *latest;
   char *directory;
   char *jobid;
   int durseconds;
} ComputeArgs;

/*
void BlockAndDoWork(const char *directory, apr_pool_t *mp, int instancenumber, int totalinstances, int secondstowait)
{
   printf("Starting compute work instance %d/%d, seconds_max %d\n", instancenumber, totalinstances, secondstowait);
   RunningWorkProcess *rwp = new RunningWorkProcess(directory, instancenumber, totalinstances, mp);
   apr_time_t start_time = apr_time_now();
   while (apr_time_sec(apr_time_now() - start_time) < secondstowait)
   {
      if (rwp->WeAreFinishedYet()) break;
      apr_sleep(1000000);
   }
   delete rwp;
}
*/

int main(int argc, const char *argv[])
{
   int nargc = argc;
   const char *const *nargv = argv;
   APRInterface apr_iface(&nargc, &nargv);
   SSLInterface ssl_iface;
   CryptInterface crypt_iface;
   SQLiteInterface sqlite_iface;
   XMLInterface xml_iface;
   static const apr_getopt_option_t opt_option[] = {
      { "host", 'n', TRUE, "hostname to connect to (broker host)" },
      { "port", 'p', TRUE, "port to connect to (broker port)" },
      { "directory", 'd', TRUE, "directory to run work in" },
      { "latest", 'z', TRUE, "latest key of the compute" },
      { "jobid", 'j', TRUE, "job id assigned by system" },
      { "duration", 't', TRUE, "maximum duration of the compute" },
      { "xmldescriptor", 'x', TRUE, "xml description of work under scrutiny" },
      { NULL, 0, 0, NULL } };
   apr_getopt_t *opt;
   int optch;
   const char *optarg = NULL;
   apr_status_t status = 0;
   apr_pool_t *mp = apr_iface.getGlobalPool();
   ComputeArgs args;
   args.hostname = apr_pstrdup(mp, DEFAULT_HOST);
   args.hostport = apr_pstrdup(mp, DEFAULT_PORT_STRING);
   args.xmlfile = apr_pstrdup(mp, "work0001.xml");
   args.latest = apr_pstrdup(mp, "compute");
   args.directory = apr_pstrdup(mp, "/dev/null");
   args.jobid = apr_pstrdup(mp, "job");
   args.durseconds = 3600 * 24;
   // 24 hours maximum time running - this should probably be changed (possible to infinite) !!!!
   apr_getopt_init(&opt, mp, nargc, nargv);
   while ((status = apr_getopt_long(opt, opt_option, &optch, &optarg)) == APR_SUCCESS)
   {
      switch (optch)
      {
         case 'n': args.hostname = apr_pstrdup(mp, optarg); break;
         case 'p': args.hostport = apr_pstrdup(mp, optarg); break;
         case 'x': args.xmlfile = apr_pstrdup(mp, optarg); break;
         case 'z': args.latest = apr_pstrdup(mp, optarg); break;
         case 'd': args.directory = apr_pstrdup(mp, optarg); break;
         case 'j': args.jobid = apr_pstrdup(mp, optarg); break;
         case 't': args.durseconds = atoi(apr_pstrdup(mp, optarg)); break;
         default:
         {
            const apr_getopt_option_t *options = opt_option;
            int count = 0;
            printf("Help for client program:\n");
            printf("   progname");
            while (options[count].optch != '\0')
            {
               printf(" -%c", options[count].optch);
               if (options[count].has_arg)
                  printf(" <%s>", options[count].name);
               count++;
            }
            printf("\n\n");
            count = 0;
            while (options[count].optch != '\0')
            {
               printf("(-%c or --%s) - %s\n",
                  options[count].optch, options[count].name, options[count].description);
               count++;
            }
            return 0;
            break;
         }
      }
   }
   if (status != APR_EOF)
      CheckAPRError(status);
   // this should be disabled, the first argument should be null
   KeyDB *keydb = new KeyDB(NULL, COMPUTENAME, NULL, COMPUTENAME, 0);
   char *tkeyname = NULL;
   if (StringUtils::readFileToString(&tkeyname, "brokerkeyname", apr_iface.getGlobalPool()))
   {
      errprint("could not read brokerkeyname");
      return 1;
   }
   char *tkeyfilename = StringUtils::mpprintf(apr_iface.getGlobalPool(), "%s.pub", tkeyname);
   keydb->InputPublicKeyFromFile(tkeyfilename, tkeyname, NULL, apr_iface.getGlobalPool());
   char *cpubkeyfname = StringUtils::mpprintf(apr_iface.getGlobalPool(), "%s.pub", args.latest);
   char *cprvkeyfname = StringUtils::mpprintf(apr_iface.getGlobalPool(), "%s.prv", args.latest);
   keydb->InputPublicKeyFromFile(cpubkeyfname, args.latest, COMPUTENAME, apr_iface.getGlobalPool());
   keydb->InputPrivateKeyFromFile(cprvkeyfname, args.latest, COMPUTENAME, apr_iface.getGlobalPool());
   
   // Do not free these!
   keydb->GetMainKeyset(&(keydb->selfkeypair), args.latest, KEY_BITS, COMPUTENAME);
   keydb->selfkey = StringUtils::varprintf("%s", args.latest);
   
   BrokerMessage *bmsg = new BrokerMessage(mp);
   bmsg->SetRecipient(tkeyname);
   bmsg->SetSender(args.latest);
   bmsg->AddTypeHeaderNOCRLF("REGISTER_COMPUTE", args.latest);
   bmsg->AddHeaderLineNOCRLF("Duration: %d", args.durseconds);
   Broker4Client *broker4 = new Broker4Client(args.hostname, atoi(args.hostport), mp);
   if (broker4->SendMessage(bmsg, keydb) != 0)
   {
      errprint("Could not send. Do we have the brokerkey (from %s:%d)?", args.hostname, atoi(args.hostport));
      delete bmsg;
      return 1;
   }
   bmsg->Reset();
   broker4->RecvMessage(bmsg, keydb);
   // So, it has been renamed.
   if ((bmsg->GetRecipient()) && (strcmp(bmsg->GetRecipient(), args.latest) == 0))
   {
      printf("Key exchange successful\n");
   } else
   {
      errprint("key exchange failed, the broker did not send correct details");
      return 1;
   }
   apr_time_t start_time = apr_time_now();
   int knownDuration = args.durseconds;
   // Set aside ten percent for start up and potential clean up!
   delete broker4;
#define TIMECONST 0.9
   while (apr_time_sec(apr_time_now() - start_time) < ((int)(knownDuration * TIMECONST)))
   {
      bmsg->Reset();
      printf("Requesting work...\n");
      bmsg->SetRecipient(tkeyname);
      bmsg->SetSender(keydb->selfkey);
      bmsg->AddTypeHeaderNOCRLF("REQUEST_WORK", keydb->selfkey);
      bmsg->AddHeaderLineNOCRLF("Time-Remaining: %d", ((int)(knownDuration * TIMECONST)) - apr_time_sec(apr_time_now() - start_time));
      broker4 = new Broker4Client(args.hostname, atoi(args.hostport), mp);
      if (broker4->SendMessage(bmsg, keydb) != 0)
      {
         errprint("Could not send request for work");
         delete bmsg;
         delete broker4;
         return 1;
      }
      bmsg->Reset();
      broker4->RecvMessage(bmsg, keydb);
      delete broker4;
      const apr_table_t *htable = bmsg->GetHeaderTable();
      char *workdatatag = NULL;
      if (bmsg->IsType("WORK_ASSIGNMENT"))
      {
         const char *lident_str = apr_table_get(htable, "Local-Ident");
         int lident = (lident_str != NULL) ? atoi(lident_str) : 0;
         const char *linst_str = apr_table_get(htable, "Total-Instances");
         int linst = (linst_str != NULL) ? atoi(linst_str) : 0;
         if (linst < 0)
         {
            errprint("problem with total instance number - broken workstore in broker!");
            linst = 1;
         }
         const char *maxtime_str = apr_table_get(htable, "Max-Time");
         int maxtime = (maxtime_str != NULL) ? atoi(maxtime_str) : 60;
         const char *wdatatag = apr_table_get(htable, "WorkDataTag");
         if (wdatatag)
         {
            workdatatag = (char *)malloc(sizeof(char)*(strlen(bmsg->GetProcessee())+1));
            strcpy(workdatatag, wdatatag);
         }
         char *workstorename = (char *)malloc(sizeof(char)*(strlen(bmsg->GetProcessee())+1));
         strcpy(workstorename, bmsg->GetProcessee());
         bmsg->TakeAndInsertPublicKeyByName(keydb, 0, workstorename);
         bmsg->TakeAndInsertPrivateKeyByName(keydb, 1, workstorename);
         bmsg->TakeAndInsertPrivateKeyByName(keydb, 2, "out_%s", workstorename);
         bmsg->TakeAndInsertPublicKeyByName(keydb, 3, "in_%s", workstorename);
         bmsg->TakeAndCTWriteOutFileByName(4, "work.xml");
         int noofblocks = 0;
         const char *noofblocks_str = apr_table_get(htable, "Block-Count");
         if ((noofblocks_str != NULL) && ((noofblocks = atoi(noofblocks_str)) > 0))
         {
            bmsg->Reset();
            SerialBlockSink *ses = new SerialBlockSink(args.directory, mp);
            int blockno = 0;
            while (blockno < noofblocks)
            {
               bmsg->SetRecipient(tkeyname);
               bmsg->SetSender("out_%s", workstorename);
               bmsg->AddTypeHeaderNOCRLF("DOWNLOAD_DO", workstorename);
               bmsg->AddHeaderLineNOCRLF("Block-Number: %d", blockno);
               if (workdatatag)
                  bmsg->AddHeaderLineNOCRLF("WorkDataTag: %s", workdatatag);
               broker4 = new Broker4Client(args.hostname, atoi(args.hostport), mp);
               broker4->MessageExchange(bmsg, keydb);
               delete broker4;
               if (bmsg->IsPayloadCount(2))
               {
                  StringUtils::aggressiveStringSanitiser((char *)bmsg->GetPayload(0));
                  char *chash = Hash::getSHA1Hash(bmsg->GetPayload(1), bmsg->GetPayloadLength(1));
                  if (strcmp(chash, bmsg->GetPayload(0)) == 0)
                  {
                     char *outbound_workkeyname = StringUtils::varprintf("out_%s", workstorename);
                     int realsize = 0;
                     int cesize = bmsg->GetPayloadLength(1);
                     printf("CESIZE of block %d bytes, hash %s\n", cesize, chash);
                     char *hashval = NULL;
                     ses->PutNextBlockDecrypted(outbound_workkeyname, keydb, bmsg->GetPayload(1), &realsize, &cesize, &hashval);
                     printf("BLOCK %d results - decoded size: %d hash: %s\n", blockno, realsize, hashval);
                     blockno++;
                  } else errprint("BLOCK %d corrupt, trying again...", blockno);
               } else errprint("malformed download message");
               bmsg->Reset();
            }
            delete ses;
            // Need to spawn the thread of the 'happy watchdog' - sending happy messages!
            printf("Starting compute work instance %d/%d, seconds_max %d\n", lident, linst, maxtime);
            RunningWorkProcess *rwp = new RunningWorkProcess(args.directory, lident, linst, mp);
            apr_time_t start_time = apr_time_now();
            apr_off_t heartbeats = 0;
            apr_off_t fails = 0;
            while (apr_time_sec(apr_time_now() - start_time) < maxtime)
            {
               if (rwp->WeAreFinishedYet()) break;
               if ((heartbeats % 3) == 0)
               {
                  bmsg->SetRecipient(tkeyname);
                  bmsg->SetSender(keydb->selfkey);
                  bmsg->AddTypeHeaderNOCRLF("HAPPY_COMPUTE", workstorename);
                  if (workdatatag)
                     bmsg->AddHeaderLineNOCRLF("WorkDataTag: %s", workdatatag);
                  bmsg->AddHeaderLineNOCRLF("Completion: %f", apr_time_sec(apr_time_now() - start_time) / (float)(0.01*maxtime));
                  if (args.jobid) bmsg->AddVerbatimPayloadString(args.jobid);
                  broker4 = new Broker4Client(args.hostname, atoi(args.hostport), mp);
                  broker4->MessageExchange(bmsg, keydb);
                  delete broker4;
                  if (strcmp(bmsg->GetType(), "HAPPY_BROKER") == 0)
                  {
                     errprint("happy message ping: I'm happy, you're happy...");
                     fails = 0;
                  }
                  else if (strcmp(bmsg->GetType(), "WORK_CANCELLED") == 0)
                  {
                     errprint("work cancelled");
                     fails = 3; // Immediate death - doesn't matter, broker probably dead, or has disowned the compute somehow...
                  }
                  else
                  {
                     errprint("garbled happy message...");
                     fails++;
                  }
                  bmsg->Reset();
               }
               // If the HAPPY_COMPUTE, HAPPY_BROKER fails more than three times, then it's likely something is wrong and this
               // needs to be abandoned!
               if (fails >= 3) break;
               apr_sleep(1000000);
               heartbeats++;
            }
            delete rwp;
            char *inbound_pubkeyname = StringUtils::varprintf("in_%s", workstorename);
            KeyPair *ikp = keydb->GetPublicKeyByName(inbound_pubkeyname);
            XMLDescriptor *xmldescr = new XMLDescriptor();
            xmldescr->readPlain("work.xml");
            xmldescr->unlinkXMLTreeForBlockTransfer();
            int blocks = 0;
            xmldescr->generateXMLTreeForBlockTransfer(NULL, inbound_pubkeyname, ikp, args.directory, 1048576, mp, &blocks);
            char *str_xmldescr = xmldescr->writePlain();
            bmsg->SetRecipient(tkeyname);
            bmsg->SetSender(workstorename);
            bmsg->AddTypeHeaderNOCRLF("BEGIN_SESSION", workstorename);
            bmsg->AddHeaderLineNOCRLF("Session-Type: RESULTS");
            bmsg->AddHeaderLineNOCRLF("Local-Ident: %d", lident);
            bmsg->AddVerbatimPayloadString(str_xmldescr);
            broker4 = new Broker4Client(args.hostname, atoi(args.hostport), mp);
            broker4->MessageExchange(bmsg, keydb);
            delete broker4;
            free(str_xmldescr);
            delete xmldescr;
            if (bmsg->IsPayloadCount(2))
            {
               StringUtils::aggressiveStringSanitiser((char *)bmsg->GetPayload(0));
               keydb->AddPublicKey(bmsg->GetPayload(0), "RSA", bmsg->GetPayload(1), NULL);
            }
            char *upload_session = (char *)malloc(sizeof(char)*(strlen(bmsg->GetPayload(0))+1));
            strcpy(upload_session, bmsg->GetPayload(0));
            int blocks_remaining = 1;
            while (blocks_remaining > 0)
            {
               int blockno = 0;
               SerialBlockSource *sbs = new SerialBlockSource(args.directory, mp);
               int rsize = STANDARD_BLOCK;
               int blockcomplete = 0;
               while (rsize == STANDARD_BLOCK)
               {
                  rsize = STANDARD_BLOCK;
                  int cesize = 0;
                  // outbound from here is inbound to the origin
                  char *outbound_keyname = StringUtils::varprintf("in_%s", workstorename);
                  char *ceblock = sbs->GetNextBlockEncrypted(outbound_keyname, keydb, &rsize, &cesize, NULL);
                  printf("CESIZE block %d is %d\n", blockno, cesize);
                  free(outbound_keyname);
                  bmsg->Reset();
                  bmsg->SetRecipient(upload_session);
                  bmsg->SetSender(workstorename);
                  bmsg->AddTypeHeaderNOCRLF("UPLOAD_OKAY", upload_session);
                  bmsg->AddHeaderLineNOCRLF("Block-Number: %d", blockno);
                  bmsg->Print();
                  broker4 = new Broker4Client(args.hostname, atoi(args.hostport), mp);
                  broker4->MessageExchange(bmsg, keydb);
                  delete broker4;
                  bmsg->Print();
                  if ((bmsg->IsType("FILECOMPLETE")) && (strcmp(bmsg->GetProcessee(), "N") == 0))
                  {
                     bmsg->Reset();
                     bmsg->SetRecipient(upload_session);
                     bmsg->SetSender(workstorename);
                     bmsg->AddTypeHeaderNOCRLF("UPLOAD_DO", upload_session);
                     bmsg->AddHeaderLineNOCRLF("Block-Number: %d", blockno);
                     char *phash = Hash::getSHA1Hash(ceblock, cesize);
                     bmsg->AddVerbatimPayloadString(phash);
                     bmsg->AddPayloadBuffer(ceblock, cesize);
                     broker4 = new Broker4Client(args.hostname, atoi(args.hostport), mp);
                     broker4->MessageExchange(bmsg, keydb);
                     delete broker4;
                     bmsg->Print();
                  } else
                  {
                     free(ceblock);
                     blockcomplete++;
                  }
                  blockno++;
               }
               blocks_remaining = blockno - blockcomplete;
               delete sbs;
            }
            if (blocks_remaining == 0)
            {
               bmsg->Reset();
               bmsg->SetRecipient(upload_session);
               bmsg->SetSender(workstorename);
               bmsg->AddTypeHeaderNOCRLF("SET_USABLE", upload_session);
               broker4 = new Broker4Client(args.hostname, atoi(args.hostport), mp);
               broker4->MessageExchange(bmsg, keydb);
               delete broker4;
               bmsg->Print();
            }
         } else errprint("No block count specified");
      } else if (strcmp(bmsg->GetType(), "WORKUNAVAILABLE") == 0)
      {
         printf("No work to do!\n");
         apr_sleep(1000);
         delete keydb;
         delete bmsg;
         return 0;
      } else break;
   }
   delete bmsg;
   delete keydb;
   return 0;
}
